"""Contains various functions for the mathematics of quantum computing."""

from numpy import matrix, ndarray, identity, kron
from math import sqrt, sin, cos, log
from decimal import Decimal            
from complexdecimal import ComplexDecimal

def pure(state, base):
    # Inappropriately named; to be phased out
    """Generates a NumPy vector for the given basis state in the given base.
    Note that the possible range runs from 0 to base-1, in accordance with
    general computing principles."""
    if state > base - 1:
        raise NameError("No such state in the given base! Off-by-one?")
    # Build a column vector of dimension "base" with 1 at index "state".
    vector = matrix([[ComplexDecimal(0,0)]]*state + [[ComplexDecimal(1,0)]] + \
                    [[ComplexDecimal(0,0)]]*(base-state-1))
    # Return this vector.
    return vector

basis = pure

def superposition(values,base,*weights):
    """superpositon(values,base[,weights]) creates a state vector of dimension
    'base' in a superposition of the states (descibed as integers) in 'values',
    and, if 'weights' is provided, weights them accordingly."""
    if not weights:
        coefficient = ComplexDecimal(
            Decimal.sqrt(Decimal(1)/Decimal(len(values))),0)
        weights = [0]*base
        for v in values:
            weights[v] = coefficient
    result = [0]*base
    for v in values:
        result[v] = weights[v]
    return matrix(result).transpose()

def entangle(left, rights):
    """entangled(left,rights) takes a given state vector 'left' and entangles
    it with the state vectors of 'rights'."""
    result = matrix([[0]]*(left.size*rights[0].size))
    for index in range(left.size):
        temp = left.copy()
        temp.fill(0)
        temp[index] = left[index]
        result = result+kron(temp,rights[index])
        # print kron(temp,rights[index])
    return result

def braket(state):
    """braket(state) onverts a bra to a ket, and vice versa. Identical to
    state.H"""
    return state.H

def density(states, probabilities):
    """density(states) calculates the density operator p for a given set of
    'states' with respective 'probabilities'."""
    p = states[0]*states[0].H*probabilities[0]
    for (index, state) in enumerate(states[1:]):
        p += state*state.H*probabilities[index]
    return p

def convert(decimal, base):
    """Converts a decimal number to an arbitrary numeric base."""
    result = ""
    while decimal > 0:
        decimal = decimal/base
        result = str(decimal%base) + result
    return result

def smeasure(n):
    """smeasure(system) measures a quantum system, resulting in its collapse.
    Please note that, if the system in question is a composite, it returns the
    state of the system in its entirety. Single-system measurements in a
    composite system should be performed via other means."""
    if sysval(n) == False:
        raise NameError("This vector is not normalized!")
    scale = [abs(n.flatten()[0])**2]
    for index in range(1,n.size):
        scale.append(abs(n.flat[index])**2 + scale[index-1])
    p = 0
    rand = random.random()
    for r in range(len(scale)):
        if rand > p and rand <= scale[r]:
            measured = [0]*r + [1] + [0]*(n.size-r-1)
            return matrix(measured).T
        else:
            p = scale[r]

def norm(vector):
    """norm(vector) returns the vector norm of a given vector."""
    result = 0
    for elem in vector.flat:
        result += abs(elem)**2
    return result

def pseparate(state, base=2, depth=None):
    # This isn't finished!
    """pseparate(state, base[, depth]) performs a "partial" separation of a
    system; that is, it returns a set of state vectors (either a total
    decomposition, or a decomposition to a supplied depth) that describes the
    statistical behavior of single systems within a composite. However, this
    state does not contain any data on entanglement, and is thus of very limited
    utility."""
    if not depth:
        depth = log(float(base),float(state.size))
    result = []
    c = 0
    while c < depth:
        sections = []
        for n in range(0,base,state.size/base):
            sections.append(state[n:n+state.size/base])
        rlist = []
        for s in sections:
            element = 0
            element = sqrt(sum([abs(n)**2 for n in s]))
            rlist += element
        result += matrix(rlist).T
    return result

def entangled(state, base):
    """This is a function that determines whether or not a composite state is
    entangled, given the number of involved systems and the
    number of basis states. We may write a more general function that takes
    only the number of bases and the vector for academic purposes, and
    perhaps one that is binary-only for speed. See 'Entanglement' in the
    presentation of findings in the thesis document for a description of the
    derivation of this algorithm."""
    # print state
    if state.size%base != 0:
        raise NameError("Error: mismatched states; state.size%base != 0")
    # if state.size == base or state.sum()%1 == 0 and not isinstance(state,ComplexDecimal):
        # mod 1 because the vector could be non-normalized... It shouldn't be,
        # in practice, but debugging with normalized vectors is a pain.
        # return False
    # Now to execute the core of the algorithm. The coefficients and basis
    # states are already sorted, by the properties of numerical notation.
    psize = base
    isentangled = True # We don't actually know, but this makes the loop run.
    while psize < state.size and isentangled:
        partitions = []
        for k in range(0,state.size,psize):
            partitions.append(state[k:k+psize])
        # I figured out how to do this! Hooray!
        partitionsd = []
        for group in partitions:
            s = Decimal('0')
            for n in group.flat:
                s += abs(n)**2
            k = Decimal.sqrt(s)
            separated = []
            for n in group.flat:
                if (k.real, k.imag) == (0, 0):
                    if n == 0: # This is debatable... It needs further study.
                        separated.append(Decimal('NaN'))
                    else:
                        separated.append(Decimal('inf'))
                else:
                    separated.append(n/k)
            partitionsd.append(separated)
        for element in range(base):
            # Rounding errors!
            c = partitionsd[0][element]
            equality = True
            for group in partitionsd[1:]:
                if c != group[element] and group[element] != Decimal('Nan'):
                    equality = False
                # return (c,group[element])
            if equality:
                isentangled = False
            else:
                psize *= base
    # psize is also the maximum depth at which the test has been performed, so
    # we return it, because that's where the breakpoint (if there is one) is.
    return (isentangled, psize)

# Transitional; I would like to replaced entangled() with the less-ambiguous
# isentangled() at some point.
# isentangled = entangled

def separate(state, base):
    """separate(state,base) separates a separable state. It operates on a
    similar principle to entangled() -- association to indices, basewise
    deconstruction, etc. Please note that there is the occasional composite
    state that can be described as more than one tensor product -- in this case,
    one of these is returned. This behavior is subject to change."""
    # Optimizing...
    if state.size%base != 0:
        raise NameError("Mismatched arguments! state.size%base != 0")
    if state.size == base:
        return state
    if state.sum() == 1:
        # This is just reversing the basis-state index calculaton, as any state
        # vector with a single entry of '1' (which is the only way any
        # normalized vector of dimension > 1 [which is kind of a ridiculous
        # thing to do] could have a coefficient sum of 1.)
        index = [n for n in state.flat].index(1)
        right = index%base
        left = (index-right)/base
        # print left, right
        rightvector = matrix([[0]]*right+[[1]]+[[0]]*(base-right-1))
        leftvector = matrix([[0]]*left+[[1]]+[[0]]*(state.size/base-left-1))
        return (leftvector, rightvector)
    
       # First, we partition in groups of length 'state.size'/'base', creating
       # collections of coefficients corresponding to a given left-hand
       # basis state vector.
    partitions = []
    step = state.size/base
    for p in range(0,state.size,step):
        partitions.append(state[p:p+step])
    # Now we calculate the left-hand side.
    left = []
    for group in partitions:
        # We use Decimals.
        summation = Decimal('0')
        for entry in group.flat:
            summation += abs(entry)**2
        left.append(summation)
    # Now we calculate the right side. This, of course, doesn't work if the
    # state is non-separable. By the properties of separable states, we know
    # that we can factor out any member of 'left' from its corresponding
    # partition to get the left-hand vector. We choose to use left[0] and
    # partitions[0], because they should always be present -- unless it equals
    # 0. If this is the case, we know that there is at least one other,
    # non-zero, entry. We iterate through 'left' to find it, so we can avoid
    # dividing by 0.
    index = 0
    while left[index] == 0:
        index += 1
    right = []
    for entry in partitions[index].flat:
        right.append(entry/left[index])
    # Hooray! We've finished.
    return (matrix(left).T, matrix(right).T)


def sysval(vector):
    """sysval(vector) checks to see if a vector can described a quantum system
    -- that is, it checks to see if it is, in fact, a unit vector."""
    # Do we need this as a separate function at all? It seems simple enough to
    # simply write in to any code that needs it.
    r = 0
    for e in vector.flat:
        r += abs(e)**2
    if r == 1:
        return True
    else:
        return False

def gval(m,base):
    """gval(matrix,base)function validates a matrix to ensure that it is
    unitary, and can consequentially be used as a quantum gate. It also checks
    to see if it is a valid quantum gate for the given numeric base."""
    if (m*m.H == matrix(identity(m.shape[0]))).all() and \
           [d%base for d in m.shape] == [0]*len(m.shape):
        return True
    else:
        return False

def pad(m,target,syslength):
    """pad(matrix,target,syslength) "pads" a matrix with identity matrices so
    that it will only act upon the system at index 'target' a composite system
    of dimension 'syslength'."""
    # iml is the pad on the left side, while imr is that on the right.
    iml = matrix(identity(m.shape[0]**target))
    imr = matrix(identity(m.shape[0]**(syslength-target-1)))
    result = kron(kron(iml,m),imr)
    return result

def controlled(m,controlcount):
    """controlled(matrix,controlcount) creates a controlled-gate matrix from a given matrix 'matrix', with 'controlcount' controls."""
    # This needs to be modified to accept varying types...
    control = identity(m.shape[0]**(controlcount+1),object)
    control[control.shape[0]-m.shape[0]:control.shape[0], \
            control.shape[1]-m.shape[1]:control.shape[1]] = m
    return control

def syscontrol(upsilon,target,control,syslength):
    """syscontrol(upsilon,target,system) creates a controlled-
    upsilon operator acting upon index 'target' of a system of length
    'syslength', with the control at index 'control' in the same. The process
    by which this occurs is described in the thesis document."""
    # First, we create the identity buffer for [upsilon].
    # -1 is the correction for from-zero indexing.
    epsilon = pad(epsilon,target-control)
    # Clearing unneeded variables.
    del(upsilon)
    # Now we create the controlled-epsilon operator.
    cepsilon = controlled(epsilon,1)
    # We can now also clear epsilon. This is really just housekeeping.
    del(epsilon)
    # Now to finish padding.
    left = matrix(identity(base**(control)))
    right = matrix(identity(base**(syslength-target-1)))
    return kron(kron(left,cepsilon),right)
